import os

from langchain.agents import ConversationalAgent, AgentExecutor
from langchain.memory import ConversationBufferMemory, ConversationSummaryMemory, ConversationBufferWindowMemory, \
    ConversationSummaryBufferMemory
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0
)


@tool(
    description="查询指定城市的天气数据,参数为city,中文城市名,字符串,示例:北京,上海,返回值:天气数据"
)
def WeatherTool(city: str) -> str:
    weather_data = {
        "北京": {"今天": "15℃", "明天": "18℃", "后天": "20℃"},
        "上海": {"今天": "25℃", "明天": "28℃", "后天": "30℃"},
        "郑州": {"今天": "35℃", "明天": "38℃", "后天": "30℃"},
    }

    return f"{city}今天:{weather_data.get(city).get("今天")},明天:{weather_data.get(city).get("明天")}"


# 12℃,20℃

@tool(description="计算两个温度的差值,参数为两个温度字符串(用逗号分割,比如:12℃,20℃),返回值:温度差")
def TemlDiffTool(teps: str) -> str:
    try:
        tep1, tep2 = teps.split(",")
        t1 = int(tep1.replace("℃", ""))
        t2 = int(tep2.replace("℃", ""))
        return f"温度差:{abs(t1 - t2)}℃"
    except Exception as e:
        return "计算失败,输入格式错误"


tools = [WeatherTool, TemlDiffTool]

# 定义记忆组件[存储,添加,更新]
# memory = ConversationBufferMemory(
#     memory_key="chat_history",
#     return_messages=True
# )

# 定义记忆组件(压缩) 100
memory = ConversationSummaryMemory(
    memory_key="chat_history",
    return_messages=True,
    llm=llm,
    max_token_limit=100

)

# memory = ConversationSummaryBufferMemory(
#     memory_key="chat_history",
#     return_messages=True,
#     llm=llm,
#     max_token_limit=100
# )

# memory = ConversationBufferWindowMemory(
#     k=3,
#     memory_key="chat_history",
#     return_messages=True
# )

prompt_template = """

你是一个智能天气助手,你拥有以下工具:
{tools},

历史对话:
{chat_history}

使用规则:
1:先理解问题:无需工具则直接回答,需要工具则调用工具(根据工具的描述匹配)
2:调用工具必须用格式：```json{{"name":"工具名","parameters":{{"参数名":"参数值"}}}}```（参数需齐全）；
3:拿到工具结果后, 判断是否继续调用工具, 如果需要继续调用工具,则调用工具,否则整理,返回结果;
4:如果问指定城市的温度差,需要先获取指定城市的温度,再计算温度差


用户的问题:{input}

"""
prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["tools", "input", "chat_history"],
    partial_variables={
        "tools": "\n".join([f"{t.name}:{t.description}" for t in tools])
    }
)

agent = ConversationalAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True,
    memory=memory
)

executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,
    memory=memory
)

if __name__ == '__main__':

    while True:
        query = input("请输入问题:")
        res = executor.invoke({
            "input": query
        })
        print(res)
        print("当前记忆:", memory.buffer)

    # print("第一轮记忆:", memory.buffer)
    # query1 = "北京今天的温度是多少?"
    # res1 = executor.invoke({
    #     "input": query1
    # })
    #
    # print(res1)
    # print("第二轮记忆:", memory.buffer)
    # query2 = "那明天呢?"
    # res2 = executor.invoke({
    #     "input": query2
    # })
    #
    # print(res2)
    #
    # print("第三轮记忆:", memory.buffer)
    # query3 = "上海今天的天气怎么样?"
    # res3 = executor.invoke({
    #     "input": query3
    # })
    #
    # print(res3)
    #
    # print("第四轮记忆:", memory.buffer)
    # query4 = "两地今天的温差是多少"
    # res4 = executor.invoke({
    #     "input": query4
    # })
    #
    # print(res4)
